import os
import glob
import random
from PIL import Image


if __name__ == '__main__':
    split_rate = 0.1      # 验证集划分比例
    resize_image = 224    # 图片缩放尺寸
    file_path = '.\\img'  # 原始数据路径

    # 获取所有类别目录
    dirs = glob.glob(os.path.join(file_path, '*'))
    dirs = [d for d in dirs if os.path.isdir(d)]

    print(f"共发现 {len(dirs)} 个类别: {dirs}")

    for class_dir in dirs:
        # 提取类别名称
        class_name = os.path.basename(class_dir)
        
        # 创建训练集和测试集目录
        os.makedirs(f"data_set\\train\\{class_name}", exist_ok=True)
        os.makedirs(f"data_set\\test\\{class_name}", exist_ok=True)

        # 获取所有图片文件
        files = glob.glob(os.path.join(class_dir, '*.jpg'))
        files += glob.glob(os.path.join(class_dir, '*.jpeg'))
        files += glob.glob(os.path.join(class_dir, '*.png'))
        
        random.shuffle(files)  # 打乱文件顺序
        total = len(files)     # 类别总数量
        split = int(total * split_rate)  # 验证集数量

        # 遍历处理每个文件
        for idx, file_path in enumerate(files):
            # 打开并调整图片尺寸
            img = Image.open(file_path).convert('RGB')
            img = img.resize((resize_image, resize_image), Image.Resampling.LANCZOS)
            
            # 生成连续编号
            if idx < split:  # 测试集文件
                number = (total - split) + idx + 1  # 从总训练数+1开始编号
                save_dir = f"data_set\\test\\{class_name}"
            else:           # 训练集文件
                number = idx - split + 1  # 从1开始编号
                save_dir = f"data_set\\train\\{class_name}"
            
            # 格式化文件名：类别+四位编号
            new_name = f"{class_name}{number:04d}.png"
            img.save(os.path.join(save_dir, new_name))

    # 统计最终数据量
    train_files = glob.glob(os.path.join('data_set', 'train', '*', '*.png'))
    test_files = glob.glob(os.path.join('data_set', 'test', '*', '*.png'))
    
    print(f"训练集共 {len(train_files)} 张图片")
    print(f"验证集共 {len(test_files)} 张图片")